import cv2
import numpy as np
import logging


def train_bg_subtractor(inst, cap, num=500):
    '''
        BG減法器需要處理一些幀來開始給出結果
    '''
    print('Training BG Subtractor...')
    i = 0
    while cap.isOpened():
        ret, frame = cap.read()
        inst.apply(frame, None, 0.001)
        i += 1
        if i >= num:
            return cap


# 創造MOG BG減法器,緩存500幀和陰影剔除
bg_subtractor = cv2.createBackgroundSubtractorMOG2(history=500, detectShadows=True)
# 設置圖像源
path = '../../../../large_data/video/traffic/traffic02.mp4'
# cap = cv2.VideoCapture('input.mp4')
cap = cv2.VideoCapture(path)
# 跳過500幀訓練bg減法器
train_bg_subtractor(bg_subtractor, cap, num=500)


def filter_mask(self, img, a=None):
    # 這個過濾器是根據視覺測試手工挑選的
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))


    # 填補任何小洞
    closing = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
    # 去除噪音
    opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)

    # 膨脹以合併相鄰的小塊
    dilation = cv2.dilate(opening, kernel, iterations=2)

    return dilation


# 得到邊界框的中心
def get_centroid(x, y, w, h):
    x1 = int(w / 2)
    y1 = int(h / 2)

    cx = x + x1
    cy = y + y1

    return (cx, cy)


def distance(pt1, pt2, wx, wy):
    x1, y1 = pt1
    x2, y2 = pt2
    return np.sqrt((x1 - x2)**2 * wx + (y1 - y2)**2 * wy)


def save_frame(img, path, flip=True):
    img = cv2.flip(img, 1)
    cv2.imwrite(path, img)


def detect_vehicles(fg_mask, min_contour_width=35, min_contour_height=35):
    matches = []
    # 發現外部輪廓
    contours, hierarchy = cv2.findContours(
        fg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)

    # 根據高度，寬度進行過濾
    for (i, contour) in enumerate(contours):
        # 找到其中一條輪廓,進行高和寬的過濾
        (x, y, w, h) = cv2.boundingRect(contour)
        contour_valid = (w >= min_contour_width) and (
                h >= min_contour_height)

        if not contour_valid:
            # 不滿足過濾掉
            continue

        # 得到邊界框的中心
        centroid = get_centroid(x, y, w, h)
        # 將輪廓加入數組中
        matches.append(((x, y, w, h), centroid))

    return matches


class PipelineRunner(object):
    '''
       非常簡單的pipline。
       只需按順序運行處理程序，將上下文從一個傳遞到另一個。
       您還可以爲處理器設置日誌級別。
    '''



def __init__(self, pipeline=None, log_level=logging.DEBUG):
    self.pipeline = pipeline or []
    self.context = {}
    self.log = logging.getLogger(self.__class__.__name__)
    self.log.setLevel(log_level)
    self.log_level = log_level
    self.set_log_level()



def set_context(self, data):
    self.context = data



def add(self, processor):
    if not isinstance(processor, PipelineProcessor):
        raise Exception(
            'Processor should be an isinstance of PipelineProcessor.')
    processor.log.setLevel(self.log_level)
    self.pipeline.append(processor)



def remove(self, name):
    for i, p in enumerate(self.pipeline):
        if p.__class__.__name__ == name:
            del self.pipeline[i]
            return True
    return False



def set_log_level(self):
    for p in self.pipeline:
        p.log.setLevel(self.log_level)



def run(self):
    for p in self.pipeline:
        self.context = p(self.context)


    self.log.debug("Frame #%d processed.", self.context['frame_number'])

    return self.context



class PipelineProcessor(object):
    '''
        處理器的基類。
    '''



    def __init__(self):
        self.log = logging.getLogger(self.__class__.__name__)


class ContourDetection(PipelineProcessor):
    '''
        檢測移動物體。
        這個處理器的目的是減去背景，獲得移動物體
        並使用cv2.findContours方法檢測它們，然後過濾掉寬度和高度不符合的。
        bg_subtractor  - 背景減法器isinstance。
        min_contour_width  - 最小邊界矩形寬度。
        min_contour_height  - 最小邊界矩形高度。
        save_image  - 如果爲True，則將檢測到的對象掩碼保存到文件中。
        image_dir  - 保存圖像的位置（必須存在）。
    '''

    def __init__(self, bg_subtractor, min_contour_width=35, min_contour_height=35, save_image=False,
                 image_dir='images'):
        super(ContourDetection, self).__init__()

        self.bg_subtractor = bg_subtractor
        self.min_contour_width = min_contour_width
        self.min_contour_height = min_contour_height
        self.save_image = save_image
        self.image_dir = image_dir

    def filter_mask(self, img, a=None):
        # 這個過濾器是根據視覺測試手工挑選的
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))

        # 閉運算，填補任何小洞
        closing = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
        # 開運算，去除噪音
        opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)

        # 膨脹以合併相鄰的小塊
        dilation = cv2.dilate(opening, kernel, iterations=2)

        return dilation

    def detect_vehicles(self, fg_mask, context):

        matches = []

        # 發現外部輪廓
        contours, hierarchy = cv2.findContours(fg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)
        for (i, contour) in enumerate(contours):
            # 找到其中一條輪廓,進行高和寬的過濾
            (x, y, w, h) = cv2.boundingRect(contour)
            contour_valid = (w >= self.min_contour_width) and (
                    h >= self.min_contour_height)

            if not contour_valid:
                # 不滿足過濾掉
                continue
            # 得到邊界框的中心
            # centroid = utils.get_centroid(x, y, w, h)
            centroid = get_centroid(x, y, w, h)
            # 將輪廓加入數組中
            matches.append(((x, y, w, h), centroid))

        return matches

    def __call__(self, context):
        frame = context['frame'].copy()
        frame_number = context['frame_number']

        fg_mask = self.bg_subtractor.apply(frame, None, 0.001)
        # 只是閾值化
        fg_mask[fg_mask < 240] = 0
        fg_mask = self.filter_mask(fg_mask, frame_number)
        if self.save_image:
            #可選,保存過濾後的黑白圖片
            # utils.save_frame(fg_mask, self.image_dir +
            save_frame(fg_mask, self.image_dir +
                             "/mask_%04d.png" % frame_number, flip=False)
        #進行目標檢測
        context['objects'] = self.detect_vehicles(fg_mask, context)
        context['fg_mask'] = fg_mask

        return context


class VehicleCounter(PipelineProcessor):
    '''
        計算進入出口區域的車輛。
        此類的目的是基於檢測到的對象和本地緩存創建
        對象修補和計數在退出區域中定義的退出區域中輸入的對象。
        exit_masks  - 退出監測後的點集合列表。
        path_size  - 路徑中的最大點數。
        max_dst  - 兩點之間的最大距離。
    '''

    def __init__(self, exit_masks=[], path_size=10, max_dst=30, x_weight=1.0, y_weight=1.0):
        super(VehicleCounter, self).__init__()

        self.exit_masks = exit_masks

        self.vehicle_count = 0
        self.path_size = path_size
        self.pathes = []
        self.max_dst = max_dst
        self.x_weight = x_weight
        self.y_weight = y_weight

    def check_exit(self, point):
        for exit_mask in self.exit_masks:
            try:
                if exit_mask[point[1]][point[0]] == 255:
                    return True
            except:
                return True
        return False

    def __call__(self, context):
        objects = context['objects']
        context['exit_masks'] = self.exit_masks
        context['pathes'] = self.pathes
        context['vehicle_count'] = self.vehicle_count
        if not objects:
            return context

        points = np.array(objects)[:, 0:2]
        points = points.tolist()

        # 如果pathes爲空，則添加新點
        if not self.pathes:
            for match in points:
                self.pathes.append([match])

        else:
            # 根據最小距離鏈接新點與舊路徑
            # points
            new_pathes = []

            for path in self.pathes:
                _min = 999999
                _match = None
                for p in points:
                    if len(path) == 1:
                        # distance from last point to current
                        # d = utils.distance(p[0], path[-1][0])
                        d = distance(p[0], path[-1][0])
                    else:
                        # 基於前2個點預測下一個點並計算
                        # 從預測的下一個點到當前的距離
                        xn = 2 * path[-1][0][0] - path[-2][0][0]
                        yn = 2 * path[-1][0][1] - path[-2][0][1]
                        # d = utils.distance(
                        d = distance(
                            p[0], (xn, yn),
                            x_weight=self.x_weight,
                            y_weight=self.y_weight
                        )

                    if d < _min:
                        _min = d
                        _match = p

                if _match and _min <= self.max_dst:
                    points.remove(_match)
                    path.append(_match)
                    new_pathes.append(path)

                # 如果當前幀沒有匹配，則不丟棄路徑
                if _match is None:
                    new_pathes.append(path)

            self.pathes = new_pathes

            # 添加新的pathes
            if len(points):
                for p in points:
                    # 不添加已經計算的點數
                    if self.check_exit(p[1]):
                        continue
                    self.pathes.append([p])

        # 僅保存路徑中的最後N個點
        for i, _ in enumerate(self.pathes):
            self.pathes[i] = self.pathes[i][self.path_size * -1:]

        # 計算車輛和去除計數pathes：
        new_pathes = []
        for i, path in enumerate(self.pathes):
            d = path[-2:]

            if (
                    # 需要列表中的兩點來計算
                    len(d) >= 2 and
                    # 上一個點不是在出口區
                    not self.check_exit(d[0][1]) and
                    # 當前點在出口區
                    self.check_exit(d[1][1]) and
                    # path len 大於 path_size
                    self.path_size <= len(path)
            ):
                self.vehicle_count += 1
            else:
                # 禁止鏈接已經在退出區域中的路徑
                add = True
                for p in path:
                    if self.check_exit(p[1]):
                        add = False
                        break
                if add:
                    new_pathes.append(path)

        self.pathes = new_pathes

        context['pathes'] = self.pathes
        context['objects'] = objects
        context['vehicle_count'] = self.vehicle_count

        self.log.debug('#VEHICLES FOUND: %s' % self.vehicle_count)

        return context


EXIT_PTS = np.array([
    [[732, 720], [732, 590], [1280, 500], [1280, 720]],
    [[0, 400], [645, 400], [645, 0], [0, 0]]
])

base = np.zeros(SHAPE + (3,), dtype='uint8')
exit_mask = cv2.fillPoly(base, EXIT_PTS, (255, 255, 255))[:, :, 0]


class Visualizer(PipelineProcessor):

    def __init__(self, save_image=True, image_dir='images'):
        super(Visualizer, self).__init__()

        self.save_image = save_image
        self.image_dir = image_dir

    def check_exit(self, point, exit_masks=[]):
        for exit_mask in exit_masks:
            if exit_mask[point[1]][point[0]] == 255:
                return True
        return False
    # 繪製路徑
    def draw_pathes(self, img, pathes):
        if not img.any():
            return

        for i, path in enumerate(pathes):
            path = np.array(path)[:, 1].tolist()
            for point in path:
                # 畫出圓點和藍色線條實現車尾尾隨效果
                cv2.circle(img, point, 2, CAR_COLOURS[0], -1)
                cv2.polylines(img, [np.int32(path)], False, CAR_COLOURS[0], 1)

        return img
    # 繪製包裹車的盒子
    def draw_boxes(self, img, pathes, exit_masks=[]):
        for (i, match) in enumerate(pathes):

            contour, centroid = match[-1][:2]
            if self.check_exit(centroid, exit_masks):
                continue

            x, y, w, h = contour
            # 根據車的寬高，繪製出車的矩形大小
            cv2.rectangle(img, (x, y), (x + w - 1, y + h - 1),
                          BOUNDING_BOX_COLOUR, 1)
            cv2.circle(img, centroid, 2, CENTROID_COLOUR, -1)

        return img

    def draw_ui(self, img, vehicle_count, exit_masks=[]):

        # 這只是爲圖像添加不透明度的綠色遮罩
        for exit_mask in exit_masks:
            _img = np.zeros(img.shape, img.dtype)
            _img[:, :] = EXIT_COLOR
            mask = cv2.bitwise_and(_img, _img, mask=exit_mask)
            cv2.addWeighted(mask, 1, img, 1, 0, img)
        # 在圖片上面繪製計數結果
        cv2.rectangle(img, (0, 0), (img.shape[1], 50), (0, 0, 0), cv2.FILLED)
        cv2.putText(img, ("Vehicles passed: {total} ".format(total=Zvehicle_count)), (30, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 1)
        return img

    def __call__(self, context):
        frame = context['frame'].copy()
        frame_number = context['frame_number']
        pathes = context['pathes']
        exit_masks = context['exit_masks']
        vehicle_count = context['vehicle_count']

        frame = self.draw_ui(frame, vehicle_count, exit_masks)
        frame = self.draw_pathes(frame, pathes)
        frame = self.draw_boxes(frame, pathes, exit_masks)

        # utils.save_frame(frame, self.image_dir +
        save_frame(frame, self.image_dir +
                         "/processed_%04d.png" % frame_number)

        return context